import cv2 as cv
import numpy as np


class DocScanner:
    def load_image(self, file_path):
        img = cv.imread(file_path)
        if img is None:
            return None, None  # 加载失败时返回None
        
        img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        # 使用自适应阈值处理不同光照条件
        binary_img = cv.adaptiveThreshold(
            img_gray, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY_INV, 11, 2
        )
        # 寻找轮廓并筛选最大四边形
        contours, _ = cv.findContours(binary_img, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        max_contour = self._find_largest_quad_contour(contours)
        
        if max_contour is None:
            return img, None  # 未检测到角落时返回None
        
        corners = self.order_points(max_contour.reshape(4, 2))
        return img, corners

    def _find_largest_quad_contour(self, contours):
        largest_contour = None
        max_area = 0
        for cnt in contours:
            perimeter = cv.arcLength(cnt, True)
            approx = cv.approxPolyDP(cnt, 0.02 * perimeter, True)
            if len(approx) == 4:
                area = cv.contourArea(cnt)
                if area > max_area:
                    max_area = area
                    largest_contour = approx
        return largest_contour

    def order_points(self, pts):
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]  # 左上
        rect[2] = pts[np.argmax(s)]  # 右下
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]  # 右上
        rect[3] = pts[np.argmax(diff)]  # 左下
        return rect.astype(int)  # 转换为整数坐标

    def crop_image(self, img, corners):
        if corners is None or len(corners) != 4:
            return None
        
        (tl, tr, br, bl) = corners
        width = max(
            np.linalg.norm(tr - br),
            np.linalg.norm(tl - bl)
        )
        height = max(
            np.linalg.norm(tl - tr),
            np.linalg.norm(bl - br)
        )
        width = int(width)
        height = int(height)
        
        dst_points = np.array([
            [0, 0],
            [width - 1, 0],
            [width - 1, height - 1],
            [0, height - 1]
        ], dtype="float32")
        
        M = cv.getPerspectiveTransform(corners.astype(float), dst_points)
        return cv.warpPerspective(img, M, (width, height))